import json
import os

import requests
import streamlit as st
from dotenv import load_dotenv

load_dotenv()


def get_response(st, model):
    messages = st.session_state.messages
    data = json.dumps({"model": model, "messages": messages})
    if model == "gpt-3.5-turbo":
        URL = os.environ.get("GPT_3.5_TURBO_URL")
        API_KEY = os.environ.get("GPT_3.5_TURBO_KEY")
    elif model == "kimi":
        URL = os.environ.get("KIMI_URL")
        API_KEY = os.environ.get("KIMI_KEY")
    elif model == "glm4":
        URL = os.environ.get("GLM_URL")
        API_KEY = os.environ.get("GLM_KEY")
    elif model == "step":
        URL = os.environ.get("STEP_CHAT_URL")
        API_KEY = os.environ.get("STEP_CHAT_KEY")
    elif model == "emohaa":
        URL = os.environ.get("EMOHAA_URL")
        API_KEY = os.environ.get("EMOHAA_KEY")
    elif model == "qwen":
        URL = os.environ.get("QWEN_URL")
        API_KEY = os.environ.get("QWEN_KEY")
    elif model == "hailuo":
        URL = os.environ.get("HAILUO_URL")
        API_KEY = os.environ.get("HAILUO_KEY")
    elif model == "spark":
        URL = os.environ.get("SPARK_URL")
        API_KEY = os.environ.get("SPARK_KEY")
    elif model == "deepseek":
        URL = os.environ.get("DEEPSEEK_URL")
        API_KEY = os.environ.get("DEEPSEEK_KEY")
    else:
        return "Error: Invalid model"

    try:
        headers = {
            "Authorization": f"Bearer {API_KEY}",
            "Content-Type": "application/json"
        }
        rs = requests.post(URL, headers=headers, data=data, stream=False)
        if rs.status_code == 200:
            return rs.json()["choices"][0]["message"]["content"]
        else:
            return f"Error: {rs.status_code}"
    except Exception as e:
        print(e)
        return f"Error: {e}"


st.title(os.environ.get("TITLE_NAME"))
if "messages" not in st.session_state:
    st.session_state.messages = []
# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])
genre = st.sidebar.selectbox(
    "请选择你的聊天模型:",
    ('gpt-3.5-turbo', 'kimi', 'glm4', 'step', 'emohaa', 'qwen', 'hailuo', 'spark', 'deepseek'))
# React to user input
if prompt := st.chat_input("想要聊点什么?"):
    # Display user message in chat message container
    st.chat_message("user").markdown(prompt)
    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": prompt})
    response = get_response(st, genre)
    # Display assistant response in chat message container
    with st.chat_message("assistant"):
        st.markdown(response)
    # Add assistant response to chat history
    st.session_state.messages.append({"role": "assistant", "content": response})
if st.sidebar.button("清空聊天记录"):
    st.session_state.messages.clear()
    st.rerun()
st.sidebar.markdown("------------------------------------------")
st.sidebar.markdown("Powered by [Johnson](https://github.com/SchrodingerFish)")
